package com.chen.authority.common.aspect;

import cn.hutool.core.text.StrPool;
import cn.hutool.extra.servlet.ServletUtil;
import com.chen.authority.common.annotation.RateLimiter;
import com.chen.authority.common.util.RequestUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ValueOperations;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * 限流切面
 *
 * @author chengy
 * @date 2024/06/14
 */
@Slf4j
@Aspect
@Component
@RequiredArgsConstructor
public class RateLimiterAspect {

    private final RedisTemplate<String, Object> redisTemplate;

    /**
     * 锁住时的key前缀
     */
    public static final String LOCK_PREFIX = "LOCK_";

    /**
     * 统计次数时的key前缀
     */
    public static final String COUNT_PREFIX = "COUNT_";

    /**
     * 限流前置处理
     *
     * @param joinPoint   切点
     * @param rateLimiter 限流注解
     */
    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint joinPoint, RateLimiter rateLimiter) {
        HttpServletRequest request = RequestUtil.getRequest();
        if (RequestUtil.isReptile(request)) {
            throw new RuntimeException("系统监测为恶意爬虫程序，禁止访问");
        }
        String ip = ServletUtil.getClientIP(request);
        String uri = request.getRequestURI();
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        String classPath = joinPoint.getTarget().getClass().getName();
        String methodName = signature.getMethod().getName();
        String methodPath = classPath + StrPool.DOT + methodName;
        if (needForbidden(ip, uri, methodPath, rateLimiter)) {
            throw new RuntimeException("您的操作过快，已被系统限流【" + rateLimiter.forbiddenSecond() + "】秒内无法进行此操作");
        }
    }

    /**
     * 判断是否需要禁用
     *
     * @param ip          请求ip
     * @param uri         请求接口地址
     * @param methodPath  方法路径
     * @param rateLimiter 限流注解
     * @return 是否需要禁用
     */
    private boolean needForbidden(String ip, String uri, String methodPath, RateLimiter rateLimiter) {
        String lockKey = LOCK_PREFIX + ip + StrPool.UNDERLINE + uri;
        Boolean hasKey = redisTemplate.hasKey(lockKey);
        // 如果存在，则已被禁用
        if (Objects.nonNull(hasKey) && hasKey) {
            return true;
        }

        ValueOperations<String, Object> opsForValue = redisTemplate.opsForValue();

        // 不拼接uri是为了防止@PathVariable注解造成的请求api变化
        String countKey = COUNT_PREFIX + ip + StrPool.UNDERLINE + methodPath;
        hasKey = redisTemplate.hasKey(countKey);
        // 如果不存在，说明第一次访问则设置锁，否则递增计数
        if (Objects.isNull(hasKey) || !hasKey) {
            opsForValue.set(countKey, 1, rateLimiter.second(), TimeUnit.SECONDS);
        } else {
            Long countValue = opsForValue.increment(countKey);
            // 超出限制，则进行禁用操作，并删除计数
            if (Objects.nonNull(countValue) && countValue > rateLimiter.count()) {
                log.info("{}禁用访问{}", ip, uri);
                opsForValue.set(lockKey, 1, rateLimiter.forbiddenSecond(), TimeUnit.SECONDS);
                redisTemplate.delete(countKey);
                return true;
            }
        }

        return false;
    }

}